#include "Mesh.h"
#include "Element.h"
#include <iostream>
#include <string>
#include <cmath>
#include <fstream>
  

/// Eigen is linear library, providing sparse matrix, linear operator,
/// and linear solver.
#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>
typedef Eigen::SparseMatrix<double> SpMat;
typedef Eigen::Triplet<double> Tri;
typedef Eigen::VectorXd RHS;
typedef Eigen::VectorXd Solution;

/** 
 * Real solution. Can be used as Dirichlet boundary conditions.
 * 
 * @param _p variable.
 * 
 * @return function value.
 */
double u(const Point<2, double> &_p);

/** 
 * Source term.
 * 
 * @param _p variable.
 * 
 * @return function value.
 */
double f(const Point<2, double> &_p);


int main(int argc, char* argv[])
{
    /// Need a reference element. For this example, only one kind of
    /// element.
    P1Element<double> ref_ele;

    /// The algebraic accuracy for the quadrature on the element is 1.
    ref_ele.read_quad_info("data/triangle.tmp_geo", 1);

    Point<2, double> left_bottom({0.0, 0.0});
    Point<2, double> right_upper({1.0, 1.0});
    
    RectangleMesh<double> mesh(left_bottom, right_upper, 64, 64);

    size_t n_ele = mesh.get_n_grids();
    size_t n_dof = ref_ele.get_n_dofs();
    size_t n_total_dofs = mesh.get_n_dofs();
    
    /// Here we use Eigen library for all linear structures and
    /// operators.
    // Matrix<double> A(n_total_dofs, n_total_dofs);
    // std::valarray<double> rhs(0.0, n_total_dofs); / Linear algebra
    // interface. To Eigen.
    SpMat K(n_total_dofs, n_total_dofs);
    RHS Rhs(n_total_dofs);
	Rhs.setZero(n_total_dofs);
    std::vector<Tri> TriList(n_ele * n_dof * n_dof);
    std::vector<Tri>::iterator it = TriList.begin(); 

    /// Assembling ... For each element in the mesh.
    for (size_t ei = 0; ei < n_ele; ei++)
    {
	/// Find the structure of the current element.
	const Geometry &the_ele = mesh.get_grid(ei);

	/// Find the global indexes of all the degree of freedoms on the
	/// current element.
	const std::valarray<size_t> &dof = the_ele.get_dof();
	/// Within the indexes, we find the coordinates of the dofs
	/// from the mesh, and set to the reference element. Actually
	/// build a mapping between the global dofs and the local
	/// ones. (The indexes and coordinates of local dofs have set
	/// inside P1 element.)
	for (size_t k = 0; k < n_dof; k++)
	{
	    ref_ele.set_global_dof(k, mesh.get_point(dof[k]));
	}
	
	/// Get the quadrature information from the reference element.
	const std::valarray<Point<2, double> > &quad_pnts = ref_ele.get_quad_points();
	const std::valarray<double> &weights = ref_ele.get_quad_weights();	
	int n_quad_pnts = ref_ele.get_n_quad_points();
	double volume = ref_ele.get_volume();

	/// Begin to assemble ...
	for (size_t l = 0; l < n_quad_pnts; l++)
	{
	    /// Jacobian information for the coordinates transformation.
	    double detJ = ref_ele.global2local_Jacobi_det(quad_pnts[l]);
	    Matrix<double> invJ = ref_ele.local2global_Jacobi(quad_pnts[l]);

	    /// Build the entities of the coefficent matrix and right hand side vector. 
	    for (size_t i = 0; i < n_dof; i++)
	    {
		std::valarray<double> grad_i = invJ * ref_ele.local_basis_gradient(i, quad_pnts[l]);
		double phi_i = ref_ele.local_basis_function(i, quad_pnts[l]);
		for (size_t j = 0; j < n_dof; j++)
		{
		    std::valarray<double> grad_j = invJ * ref_ele.local_basis_gradient(j, quad_pnts[l]);
		    /// a(grad phi_i, grad phi_j)
		    double cont = (grad_i[0] * grad_j[0] + grad_i[1] * grad_j[1]) * weights[l] * detJ * volume;
		    // A(dof[i], dof[j]) += cont;
		    *it = Tri(dof[i], dof[j], cont);
		    it++;
		}
		/// rhs(i) = \int f(x_i) phi_i dx
		// rhs[dof[i]] += f(ref_ele.local2global(quad_pnts[l])) * phi_i * weights[l] * detJ * volume;
		Rhs[dof[i]] += f(ref_ele.local2global(quad_pnts[l])) * phi_i * weights[l] * detJ * volume;
	    }
	}
    }
    K.setFromTriplets(TriList.begin(), TriList.end());
    K.makeCompressed();

    const std::valarray<Point<2, double> >& global_dofs = mesh.get_dofs(); 
    for (size_t i = 0; i < n_total_dofs; i++)
    {
	int bm = global_dofs[i].get_boundary_mark();
	if (bm == 0)
	    continue;
	else if (bm == 1)
	{
	    double boundary_value = u(global_dofs[i]);
	    Rhs[i] = boundary_value * K.coeffRef(i, i);
	    for (Eigen::SparseMatrix<double>::InnerIterator it(K, i); it; ++it)
	    {
		size_t row = it.row();
		if (row == i)
		    continue;
		Rhs[row] -= K.coeffRef(i, row) * boundary_value;
		K.coeffRef(i, row) = 0.0;
		K.coeffRef(row, i) = 0.0;
	    }
	}
    };

//    std::cout << K << std::endl;
//    std::cout << Rhs << std::endl;

    Eigen::ConjugateGradient<Eigen::SparseMatrix<double> > Solver_sparse;
    Eigen::VectorXd x1_sparse;
    Solver_sparse.setTolerance(1e-15);
    Solver_sparse.compute(K);
    Solution solution= Solver_sparse.solve(Rhs);

    /// stupid error!
    double error = 0.0;
    for (size_t i = 0; i < n_total_dofs; i++)
	error += (solution[i] - u(global_dofs[i])) * (solution[i] - u(global_dofs[i]));
    error = std::sqrt(error);
    std::cout << "error: " << error << std::endl;

    /// Output the solution to file in OpenDx format. 
    std::ofstream output("u.dx");
    output << "object 1 class array type float rank 1 shape 2 item "
	   << n_total_dofs
	   << " data follows" << std::endl;
    output.setf(std::ios::fixed);
    output.precision(20);
    for (size_t i = 0; i < n_total_dofs; i++)
    {
	const Point<2, double> &pnt = mesh.get_point(i);
	output << pnt[0] << "\t" << pnt[1] << std::endl;
    }

    output << std::endl;
    output << "object 2 class array type int rank 1 shape 3 item "
	   << n_ele << " data follows" << std::endl;

    for (size_t i = 0; i < n_ele; i++)
    {
	const Geometry &the_ele = mesh.get_grid(i);
	output << the_ele.get_vertex(0) << "\t"
	       << the_ele.get_vertex(1) << "\t"
	       << the_ele.get_vertex(2) << std::endl;
    }
    output << "attribute \"element type\" string \"triangles\"" << std::endl;
    output << "attribute \"ref\" string \"positions\"" << std::endl;
    output << std::endl;
    output << "object 3 class array type float rank 0 item "
	   << n_total_dofs
	   << " data follows" << std::endl;
    for (size_t i = 0; i < n_total_dofs; i++)
	output << solution[i] << std::endl;

    output << "attribute \"dep\" string \"positions\"" << std::endl;
    output << std::endl;
    output << "object \"FEMFunction-2d\" class field" << std::endl;
    output << "component \"positions\" value 1" << std::endl;
    output << "component \"connections\" value 2" << std::endl;
    output << "component \"data\" value 3" << std::endl;
    output << "end" << std::endl;

    output.close();
    return 0;
};

double u(const Point<2, double> &_p)
{
    return std::sin(M_PI * _p[0]) * std::sin(2.0 * M_PI * _p[1]);
};

double f(const Point<2, double> &_p)
{
    return 5.0 * M_PI * M_PI * u(_p);
};


